import joblib
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report, roc_auc_score, f1_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, MinMaxScaler, StandardScaler, OneHotEncoder
from imblearn.over_sampling import RandomOverSampler, SMOTE, ADASYN
from sklearn.linear_model import LogisticRegression

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

train_data = pd.read_csv('../data/train.csv')
train_data.drop(['EmployeeNumber', 'Over18', 'StandardHours'], axis=1, inplace=True)
# Resources
X = train_data.drop('Attrition', axis=1)
# Target variable
y = train_data['Attrition']
# 打印特征和标签的形状
print(f"特征的形状{X.shape}")
# 打印样本的类别分布
print(f"样本的类别分布:\n{y.value_counts()}")
# 使用标签编码器



# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
# 打印训练集和测试集的形状
print("训练集的形状")
print(X_train.shape)
print(y_train.shape)
print(f"训练集的类别分布:\n{y_train.value_counts()}")
print("测试集的形状")
print(X_test.shape)
print(y_test.shape)
print(f"测试集的类别分布:\n{y_test.value_counts()}")

# 标签编码
label_encoders = {}

for col in X_train.select_dtypes(include='object').columns:
    le = LabelEncoder()
    X_train[col] = le.fit_transform(X_train[col])
    label_encoders[col] = le  # 保存编码器用于预测阶段
    X_test[col] = le.transform(X_test[col])  # 注意：不能 fit，只 transform

# 保存 label_encoders 到文件
joblib.dump(label_encoders, '../encoder/lo_label_encoders.pkl')

# 标准化
scaler = StandardScaler()
X_train = pd.DataFrame(scaler.fit_transform(X_train), columns=X_train.columns)
X_test = pd.DataFrame(scaler.transform(X_test), columns=X_test.columns)
# 保存 scaler
joblib.dump(scaler, '../encoder/lo_scaler.pkl')

# 训练集过采样，给少数类样本进行增加近似值样本
adasyn = ADASYN(random_state=42)
X_train, y_train = adasyn.fit_resample(X_train, y_train)
# 过采样后的训练集的形状
print(f"过采样后的训练集的形状{X_train.shape}")
# 打印样本的类别分布
print(f"过采样后的训练集的类别分布:\n{y_train.value_counts()}")

# 模型训练
lo = LogisticRegression()
lo = lo.fit(X_train, y_train)
y_pred = lo.predict(X_test)
#
print(f"预测结果的形状为:{y_pred.shape}")

confusion_matrix(y_test, y_pred)
sns.heatmap(confusion_matrix(y_test, y_pred), annot=True)
plt.title("混淆矩阵")
plt.savefig('../img/逻辑回归训练模型预测结果混淆矩阵.png', dpi=300, bbox_inches='tight')
plt.show()


# 评估模型
print("\n=== 最终模型评估 ===")
y_pred_proba = lo.predict_proba(X_test)[:, 1]
print(f"AUC: {roc_auc_score(y_test, y_pred_proba):.4f}")
print(f"F1-score: {f1_score(y_test, y_pred):.4f}")
print("\n分类报告:")
print(classification_report(y_test, y_pred))
print("开始模型保存...")
joblib.dump(lo, '../model/lo_best.pkl')

